import os
import tqdm
from param import args
import torch.nn as nn
import torch.optim as optim
from Functions import *
from sklearn.metrics import roc_auc_score
from FM import FM
from FFM import FFM
from DeepFM import DeepFM

def train():
    for epoch in range(args.epoch):
        loss_fun = nn.BCELoss()
        optimizer = optim.Adam(model.parameters(),lr=args.lr,weight_decay=1e-4)

        train_loss = []; val_loss = [];pred_y = [];true_y = []
        model.train()
        for x,y in train_data:
            pred = model(x)
            optimizer.zero_grad()
            loss = loss_fun(pred,y)
            loss.backward()
            optimizer.step()
            train_loss.append(loss.item())
        model.eval()
        with torch.no_grad():
            for x,y in val_data:
                pred = model(x)
                loss = loss_fun(pred,y)
                val_loss.append(loss.item())
                pred_y.extend(pred.tolist())
                true_y.extend(y.tolist())
        val_auc = roc_auc_score(y_true=true_y,y_score=pred_y)
        print("EPOCH %s train loss : %.5f   validation loss : %.5f   validation auc is %.5f" % (
        epoch, np.mean(train_loss), np.mean(val_loss), val_auc))

def get_FM():
    return FM(fields,args.laten_dim)
def get_FFM():
    return FFM(fields,args.laten_dim)

def get_DeepFM():
    return DeepFM(fields,args.laten_dim)

if __name__ == '__main__':
    train_data, val_data, _, fields = get_dataloader()
    model = eval('get_'+input("choose your model:"))
    train()